# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""eval script"""
import time
import argparse
from mindspore import context
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from src.data import create_testdataset
from src.networks import define_G
from src.loss import InvertibleRescalingNetLoss
from src.util import *
from src.metrics import calculate_ssim, calculate_psnr

if __name__ == '__main__':
    """eval"""
    parser = argparse.ArgumentParser(description='IIR test')

    # Test specifications
    parser.add_argument('--scale', type=int, default=2, choices=(2, 4),
                        help='Super resolution upscale factor.')
    parser.add_argument('--dataset_GT', type=str, default='/home/yeyin/IIR_msp_final/data/DIV2K_train_HR',
                        help='Path to the folder where the intended GT dataset is stored.')
    parser.add_argument('--resume_state', type=str,
                        default='/home/yeyin/IIR_msp_2021_1021/RCAN/ckpt/InvertibleRescaleNet_1400.ckpt',
                        help='Path to the checkpoint.')
    parser.add_argument("--device_id", type=int, default=0,
                        help="Device id, default: 0.")
    parser.add_argument('--device_target', type=str, default='GPU', choices=("GPU", "CPU"),
                        help="Device target, support GPU, CPU.")

    # Network specifications
    parser.add_argument("--subnet_type", type=str, default='DBNet',
                        help="Type of subnet, default: DBNet.")
    parser.add_argument("--in_nc", type=int, default=3,
                        help="Input channels, default: 3.")
    parser.add_argument("--out_nc", type=int, default=3,
                        help="Output channels, default: 3.")
    parser.add_argument("--block_num", type=int, default=[8, 8],
                        help="The block num, default: [8, 8].")
    parser.add_argument("--G_scale", type=int, default=2,
                        help="Super resolution upscale factor, default: false.")
    parser.add_argument("--G_init", type=str, default='xavier',
                        help="Init method of net, default: false.")

    # Loss specifications
    parser.add_argument("--pixel_criterion_forw", type=str, default='l2',
                        help="The pixel criterion of forward, default: l2.")
    parser.add_argument("--pixel_criterion_back", type=str, default='l1',
                        help="The pixel criterion of backward, default: l1.")
    parser.add_argument("--lambda_fit_forw", type=float, default=16,
                        help="lambda_fit_forw")
    parser.add_argument("--lambda_rec_back", type=int, default=1,
                        help="lambda_rec_back")
    parser.add_argument("--lambda_ce_forw", type=int, default=1,
                        help="lambda_ce_forw")
    parser.add_argument("--gradient_clipping", type=int, default=10,
                        help="The values of gradient_clipping.")
    parser.add_argument('--gaussian_scale', type=int, default=1,
                        help='The gaussian scale of net.')

    time_start = time.time()
    print("Start eval function!")
    args = parser.parse_args()

    # set content
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU", device_id=args.device_id, save_graphs=False)
    context.set_context(max_call_depth=10000)

    # create dataset
    val_ds = create_testdataset(dataset_path=args.dataset_GT, scale=args.scale)
    val_iter = val_ds.create_dict_iterator()

    # define network
    net = define_G(args)

    # define network with loss
    loss = InvertibleRescalingNetLoss(args, net)

    # loading resume state if exists
    if args.resume_state is not None:
        param_dict = load_checkpoint(args.resume_state)
        load_param_into_net(net, param_dict)
        print("saved model restore! " + str(args.resume_state))

    # init args
    idx = 0
    test_hr_psnr = []
    test_hr_ssim = []
    test_y_hr_psnr = []
    test_y_hr_ssim = []

    test_lr_psnr = []
    test_y_lr_psnr = []
    test_lr_ssim = []
    test_y_lr_ssim = []

    for _ in range(val_ds.get_dataset_size()):
        idx += 1
        val = next(val_iter)
        lq = ms.Tensor(val["LR"], ms.float16)
        gt = ms.Tensor(val["HR"], ms.float16)
        images = loss.test(lq, gt)
        sr_img = tensor2img(images[3])
        gt_img = tensor2img(images[0])
        gt_lr_img = tensor2img(images[1])
        lq_img = tensor2img(images[2])

        gt_img = gt_img / 255.
        sr_img = sr_img / 255.
        gt_lr_img = gt_lr_img / 255.
        lq_img = lq_img / 255.

        crop_size = args.scale
        cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :]
        cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :]
        psnr = calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255)
        ssim = calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255)
        test_hr_psnr.append(psnr)
        test_hr_ssim.append(ssim)

        lr_psnr = calculate_psnr(gt_lr_img * 255, lq_img * 255)
        lr_ssim = calculate_ssim(gt_lr_img * 255, lq_img * 255)
        test_lr_psnr.append(lr_psnr)
        test_lr_ssim.append(lr_ssim)

        avg_PSNR = sum(test_hr_psnr) / len(test_hr_psnr)
        avg_SSIM = sum(test_hr_ssim) / len(test_hr_ssim)
        lr_avg_PSNR = sum(test_lr_psnr) / len(test_lr_psnr)
        lr_avg_SSIM = sum(test_lr_ssim) / len(test_lr_ssim)

        if gt_img.shape[2] == 3:  # RGB image
            sr_img_y = bgr2ycbcr(sr_img, only_y=True)
            gt_img_y = bgr2ycbcr(gt_img, only_y=True)

            cropped_sr_img_y = sr_img_y[crop_size:-
            crop_size, crop_size:-crop_size]
            cropped_gt_img_y = gt_img_y[crop_size:-
            crop_size, crop_size:-crop_size]

            psnr_y = calculate_psnr(
                cropped_sr_img_y * 255, cropped_gt_img_y * 255)
            ssim_y = calculate_ssim(
                cropped_sr_img_y * 255, cropped_gt_img_y * 255)
            test_y_hr_psnr.append(psnr_y)
            test_y_hr_ssim.append(ssim_y)

            lr_img_y = bgr2ycbcr(lq_img, only_y=True)
            lrgt_img_y = bgr2ycbcr(gt_lr_img, only_y=True)
            psnr_y_lr = calculate_psnr(lr_img_y * 255, lrgt_img_y * 255)
            ssim_y_lr = calculate_ssim(lr_img_y * 255, lrgt_img_y * 255)
            test_y_lr_psnr.append(psnr_y_lr)
            test_y_lr_ssim.append(ssim_y_lr)

            avg_PSNR_y = sum(test_y_hr_psnr) / len(test_y_hr_psnr)
            avg_SSIM_y = sum(test_y_hr_ssim) / len(test_y_hr_ssim)
            lr_avg_PSNR_y = sum(test_y_lr_psnr) / len(test_y_lr_psnr)
            lr_avg_SSIM_y = sum(test_y_lr_ssim) / len(test_y_lr_ssim)

            print(
                '{:4d} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}. LR PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'.
                format(idx, psnr, ssim, psnr_y, ssim_y, lr_psnr, lr_ssim, psnr_y_lr, ssim_y_lr))
            print('      - avg PSNR: {:.6f} dB; avg SSIM: {:.6f}; avg lr PSNR: {:.6f}; avg lr SSIM: {:.6f}.'.
                  format(avg_PSNR, avg_SSIM, lr_avg_PSNR, lr_avg_SSIM))
            print('      - avg PSNR Y: {:.6f} dB; avg SSIM Y: {:.6f}; avg lr PSNR_Y: {:.6f};avg lr SSIM_Y: {:.6f}.'.
                  format(avg_PSNR_y, avg_SSIM_y, lr_avg_PSNR_y, lr_avg_SSIM_y))
        else:
            print('{:4d} - PSNR: {:.6f} dB; SSIM: {:.6f}. LR PSNR: {:.6f} dB; SSIM: {:.6f}.'.
                  format(idx, psnr, ssim, lr_psnr, lr_ssim))
            print('      - avg PSNR: {:.6f} dB; avg SSIM: {:.6f}; avg lr PSNR: {:.6f}; avg lr SSIM: {:.6f}.'.
                  format(avg_PSNR, avg_SSIM, lr_avg_PSNR, lr_avg_SSIM))
    print('eval_time: %f' % (time.time() - time_start))
